#!/bin/bash

# tailscale_optimizer.sh
# 用于自动监控和优化 Tailscale 性能

LOG_FILE="/var/log/tailscale_optimizer.log"
PID_FILE="/var/run/tailscale_optimizer.pid"
CPU_THRESHOLD=50  # CPU 使用率阈值
CHECK_INTERVAL=300  # 检查间隔（秒）
DERP_LATENCY_THRESHOLD=100  # DERP 延迟阈值（毫秒）
MAX_THREADS=4  # 最大线程数

# 检查必要的命令
check_requirements() {
    local required_commands=("taskset" "renice" "ionice" "bc" "jq")
    local missing_commands=()
    
    for cmd in "${required_commands[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing_commands+=("$cmd")
        fi
    done
    
    if [ ${#missing_commands[@]} -ne 0 ]; then
        echo "Missing required commands: ${missing_commands[*]}"
        echo "Please install the necessary packages:"
        echo "yum install -y libcgroup-tools bc jq util-linux psmisc"
        exit 1
    fi
    return 0  # 添加这行
}

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

debug_log() {
    echo "[DEBUG][$(date '+%Y-%m-%d %H:%M:%S')] $1" >> "$LOG_FILE"
}

check_cpu_usage() {
    local cpu_usage
    cpu_usage=$(ps aux | grep tailscaled | grep -v grep | awk '{print $3}' | cut -d. -f1)
    if [ -z "$cpu_usage" ]; then
        echo "0"
    else
        echo "$cpu_usage"
    fi
}

optimize_process() {
    local pid=$1
    
    if [ -z "$pid" ]; then
        log "No PID provided for optimization"
        return 1
    fi
    
    if ! kill -0 "$pid" 2>/dev/null; then
        log "Process $pid does not exist"
        return 1
    fi  # 这里应该是 fi 而不是 }
    
    # 设置进程优先级
    renice -n 10 -p "$pid" >/dev/null 2>&1
    
    # 设置 CPU 亲和性（限制使用前两个核心）
    taskset -pc 0-1 "$pid" >/dev/null 2>&1
    
    # 设置 IO 优先级
    ionice -c 2 -n 7 -p "$pid" >/dev/null 2>&1
    
    # 应用 cgroup 限制
    if command -v cgcreate &>/dev/null; then
        if [ ! -d "/sys/fs/cgroup/cpu/tailscale" ]; then
            cgcreate -g cpu:/tailscale
            echo 50000 > /sys/fs/cgroup/cpu/tailscale/cpu.cfs_quota_us
            echo 100000 > /sys/fs/cgroup/cpu/tailscale/cpu.cfs_period_us
        fi
        cgclassify -g cpu:/tailscale "$pid"
    fi
    
    log "Process $pid optimized"
    return 0
}

get_derp_info() {
    local derp_output
    derp_output=$(tailscale netcheck 2>&1)
    debug_log "Raw DERP output:"
    debug_log "$derp_output"
    
    # 提取所有 DERP 服务器信息
    echo "$derp_output" | while IFS= read -r line; do
        if [[ $line =~ ^[[:space:]]*-[[:space:]]+(.*)[[:space:]]+([0-9.]+)ms[[:space:]]+\((.*)\) ]]; then
            local server="${BASH_REMATCH[1]}"
            local latency="${BASH_REMATCH[2]}"
            local location="${BASH_REMATCH[3]}"
            echo "$server|$latency|$location"
        fi
    done
}

find_best_derp() {
    local derp_info
    derp_info=$(get_derp_info)
    debug_log "Processed DERP info:"
    debug_log "$derp_info"
    
    # 首先尝试找到亚洲区域的服务器
    local asia_servers
    asia_servers=$(echo "$derp_info" | grep -E "Hong Kong|Tokyo|Singapore" | sort -t'|' -k2 -n | head -1)
    
    if [ ! -z "$asia_servers" ]; then
        debug_log "Found Asia server: $asia_servers"
        echo "$asia_servers"
        return
    fi
    
    # 如果没有亚洲服务器，返回延迟最低的服务器
    local best_server
    best_server=$(echo "$derp_info" | sort -t'|' -k2 -n | head -1)
    if [ ! -z "$best_server" ]; then
        debug_log "Found best non-Asia server: $best_server"
        echo "$best_server"
    fi
}

optimize_derp() {
    log "Starting DERP optimization..."
    
    local derp_info
    derp_info=$(find_best_derp)
    
    if [ ! -z "$derp_info" ]; then
        IFS='|' read -r server latency location <<< "$derp_info"
        log "Found DERP server: $server (${location}, latency: ${latency}ms)"
        
        if [ ! -z "$latency" ] && (( $(echo "$latency < $DERP_LATENCY_THRESHOLD" | bc -l) )); then
            log "Switching to DERP server: $server"
            tailscale up --advertise-exit-node --exit-node-allow-lan-access >/dev/null 2>&1
            
            # 验证切换后的状态
            sleep 5
            local status_output
            status_output=$(tailscale status --json 2>/dev/null)
            log "Tailscale status after switch: $(echo "$status_output" | jq -r '.BackendState')"
        else
            log "Current DERP latency (${latency}ms) exceeds threshold (${DERP_LATENCY_THRESHOLD}ms)"
        fi
    else
        log "No suitable DERP server found in netcheck output"
    fi
}

optimize_connections() {
    local active_conns
    active_conns=$(tailscale status 2>/dev/null | grep -c "active")
    
    log "Current active connections: $active_conns"
    
    if [ "$active_conns" -gt 10 ]; then
        log "Too many active connections ($active_conns). Optimizing..."
        
        # 获取活跃连接的详细信息
        local active_info
        active_info=$(tailscale status | grep "active" | sort -k5 -n -r | head -5)
        log "Most active connections:"
        echo "$active_info" | while read -r line; do
            log "  $line"
        done
        
        tailscale up --accept-routes=false --shields-up --host-routes=false >/dev/null 2>&1
        log "Applied connection optimization settings"
    fi
}

check_and_restart() {
    local cpu_usage
    cpu_usage=$(check_cpu_usage)
    
    log "Current CPU usage: ${cpu_usage}%"
    
    if [ "$cpu_usage" -gt "$CPU_THRESHOLD" ]; then
        log "High CPU usage detected: ${cpu_usage}%"
        
        # 获取 Tailscale 进程 ID
        local tailscale_pid
        tailscale_pid=$(pgrep tailscaled)
        
        if [ ! -z "$tailscale_pid" ]; then
            log "Optimizing Tailscale process (PID: $tailscale_pid)"
            optimize_process "$tailscale_pid"
            
            # 检查线程数
            local thread_count
            thread_count=$(ls /proc/$tailscale_pid/task | wc -l)
            log "Current thread count: $thread_count"
        fi
        
        # 收集详细诊断信息
        {
            echo "=== $(date '+%Y-%m-%d %H:%M:%S') Detailed Diagnostic Information ==="
            echo "=== Tailscale Version ==="
            tailscale version
            echo "=== Netcheck Output ==="
            tailscale netcheck
            echo "=== Status Output ==="
            tailscale status
            echo "=== System Resources ==="
            free -m
            df -h
            echo "=== Thread Information ==="
            ps -T -p $tailscale_pid
            echo "=== Network Statistics ==="
            netstat -s | grep -E "segments|packets|retransmit|timeout"
            echo "=== Process Information ==="
            ps aux | grep tailscale | grep -v grep
            echo "========================="
        } >> "$LOG_FILE" 2>&1
        
        # 执行优化措施
        optimize_derp
        optimize_connections
        
        # 监控优化效果
        sleep 60
        cpu_usage=$(check_cpu_usage)
        if [ "$cpu_usage" -gt "$CPU_THRESHOLD" ]; then
            log "CPU usage still high (${cpu_usage}%). Restarting tailscaled..."
            systemctl restart tailscaled
            sleep 30
            
            # 重启后立即应用优化
            tailscale_pid=$(pgrep tailscaled)
            if [ ! -z "$tailscale_pid" ]; then
                optimize_process "$tailscale_pid"
            fi
            
            local new_cpu_usage
            new_cpu_usage=$(check_cpu_usage)
            log "Tailscaled restarted. New CPU usage: ${new_cpu_usage}%"
            
            # 验证服务状态
            local service_status
            service_status=$(systemctl is-active tailscaled)
            log "Tailscaled service status: $service_status"
        fi
    fi
}

setup_monitoring() {
    # 设置系统参数
    {
        sysctl -w net.ipv4.tcp_fastopen=3
        sysctl -w net.core.rmem_max=2500000
        sysctl -w net.core.wmem_max=2500000
        sysctl -w net.ipv4.tcp_keepalive_time=60
        sysctl -w net.ipv4.tcp_keepalive_intvl=10
        sysctl -w net.ipv4.tcp_keepalive_probes=6
        sysctl -w net.ipv4.tcp_fin_timeout=15
        sysctl -w net.ipv4.tcp_max_syn_backlog=4096
        sysctl -w net.core.netdev_max_backlog=4096
        
        # 优化进程调度
        sysctl -w kernel.sched_autogroup_enabled=0
        sysctl -w kernel.sched_tunable_scaling=0
        sysctl -w kernel.sched_migration_cost_ns=5000000
        sysctl -w kernel.sched_min_granularity_ns=10000000
        sysctl -w vm.swappiness=10
    } >/dev/null 2>&1
    
    # 初始化进程优化
    local tailscale_pid
    tailscale_pid=$(pgrep tailscaled)
    if [ ! -z "$tailscale_pid" ]; then
        optimize_process "$tailscale_pid"
    fi
    
    log "System parameters and process optimization completed"
}

cleanup() {
    log "Stopping Tailscale optimizer..."
    
    # 清理 cgroup 配置
    if [ -d "/sys/fs/cgroup/cpu/tailscale" ]; then
        cgdelete -g cpu:/tailscale >/dev/null 2>&1
    fi
    
    rm -f "$PID_FILE"
    exit 0
}

main() {
    # 检查必要的命令
    check_requirements
    
    # 检查是否已经在运行
    if [ -f "$PID_FILE" ]; then
        if kill -0 $(cat "$PID_FILE") 2>/dev/null; then
            log "Script is already running"
            exit 1
        else
            rm -f "$PID_FILE"
        fi
    fi
    
    # 写入 PID 文件
    echo $$ > "$PID_FILE"
    
    log "Starting Tailscale optimizer v1.0..."
    setup_monitoring
    
    # 设置 trap
    trap cleanup SIGTERM SIGINT SIGQUIT
    
    while true; do
        check_and_restart
        sleep "$CHECK_INTERVAL"
    done
}

# 启动主程序
main